package com.cetc.sdp.kmga.cs.util;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.hadoop.conf.Configuration;
import org.apache.hadoop.hbase.HBaseConfiguration;
import org.apache.hadoop.hbase.TableName;
import org.apache.hadoop.hbase.client.*;
import org.apache.hadoop.hbase.io.ImmutableBytesWritable;
import org.apache.hadoop.hbase.mapreduce.TableOutputCommitter;
import org.apache.hadoop.hbase.util.Bytes;
import org.apache.hadoop.mapreduce.*;

import java.io.IOException;
import java.util.HashMap;
import java.util.Map;

/**
 * 共享连接
 * Created by DengQiang on 2018/5/10.
 */
public class SharedMultiTableOutputFormat extends OutputFormat<ImmutableBytesWritable, Mutation> {
    /** Set this to {@link #WAL_OFF} to turn off write-ahead logging (WAL) */
    public static final String WAL_PROPERTY = "hbase.mapreduce.multitableoutputformat.wal";
    /** Property value to use write-ahead logging */
    public static final boolean WAL_ON = true;
    /** Property value to disable write-ahead logging */
    public static final boolean WAL_OFF = false;
    /**
     * Record writer for outputting to multiple HTables.
     */
    protected static class MultiTableRecordWriter extends
            RecordWriter<ImmutableBytesWritable, Mutation> {
        private static final Log LOG = LogFactory.getLog(SharedMultiTableOutputFormat.MultiTableRecordWriter.class);
        private static volatile Connection connection;
        Map<ImmutableBytesWritable, BufferedMutator> mutatorMap = new HashMap<>();
        Configuration conf;
        boolean useWriteAheadLogging;

        /**
         * @param conf
         *          HBaseConfiguration to used
         * @param useWriteAheadLogging
         *          whether to use write ahead logging. This can be turned off (
         *          <tt>false</tt>) to improve performance when bulk loading data.
         */
        public MultiTableRecordWriter(Configuration conf,
                                      boolean useWriteAheadLogging) {
            LOG.debug("Created new MultiTableRecordReader with WAL "
                    + (useWriteAheadLogging ? "on" : "off"));
            this.conf = conf;
            this.useWriteAheadLogging = useWriteAheadLogging;
        }

        /**
         * @param tableName
         *          the name of the table, as a string
         * @return the named mutator
         * @throws IOException
         *           if there is a problem opening a table
         */
        BufferedMutator getBufferedMutator(ImmutableBytesWritable tableName) throws IOException {
            if(connection == null) {
                //多个线程共享单个连接
                synchronized (MultiTableRecordWriter.class) {
                    if (connection == null) {
                        connection = ConnectionFactory.createConnection(conf);
                    }
                }
            }
            if (!mutatorMap.containsKey(tableName)) {
                LOG.debug("Opening HTable \"" + Bytes.toString(tableName.get())+ "\" for writing");

                BufferedMutator mutator =
                        connection.getBufferedMutator(TableName.valueOf(tableName.get()));
                mutatorMap.put(tableName, mutator);
            }
            return mutatorMap.get(tableName);
        }

        @Override
        public void close(TaskAttemptContext context) throws IOException {
            for (BufferedMutator mutator : mutatorMap.values()) {
                mutator.close();
            }
            //由于SparkStreaming任务长时间运行，因此不关闭连接
//            if (connection != null) {
//                connection.close();
//            }
        }

        /**
         * Writes an action (Put or Delete) to the specified table.
         *
         * @param tableName
         *          the table being updated.
         * @param action
         *          the update, either a put or a delete.
         * @throws IllegalArgumentException
         *          if the action is not a put or a delete.
         */
        @Override
        public void write(ImmutableBytesWritable tableName, Mutation action) throws IOException {
            BufferedMutator mutator = getBufferedMutator(tableName);
            // The actions are not immutable, so we defensively copy them
            if (action instanceof Put) {
                Put put = new Put((Put) action);
                put.setDurability(useWriteAheadLogging ? Durability.SYNC_WAL
                        : Durability.SKIP_WAL);
                mutator.mutate(put);
            } else if (action instanceof Delete) {
                Delete delete = new Delete((Delete) action);
                mutator.mutate(delete);
            } else
                throw new IllegalArgumentException(
                        "action must be either Delete or Put");
        }
    }

    @Override
    public void checkOutputSpecs(JobContext context) throws IOException,
            InterruptedException {
        // we can't know ahead of time if it's going to blow up when the user
        // passes a table name that doesn't exist, so nothing useful here.
    }

    @Override
    public OutputCommitter getOutputCommitter(TaskAttemptContext context)
            throws IOException, InterruptedException {
        return new TableOutputCommitter();
    }

    @Override
    public RecordWriter<ImmutableBytesWritable, Mutation> getRecordWriter(TaskAttemptContext context)
            throws IOException, InterruptedException {
        Configuration conf = context.getConfiguration();
        return new SharedMultiTableOutputFormat.MultiTableRecordWriter(HBaseConfiguration.create(conf),
                conf.getBoolean(WAL_PROPERTY, WAL_ON));
    }

}

